
#import "CantripSocket.h"
#import <fcntl.h>
#import <netdb.h>
#import <sys/socket.h>
#import <sys/time.h>
#import <sys/types.h>
#import <sys/ioctl.h>
#import <arpa/inet.h>
#import <unistd.h>

#define NetSocketInvalidPort				0
#define NetSocketMaxPendingConnections		5

#pragma mark -

static void callbackSocket( CFSocketRef inSocketRef, CFSocketCallBackType inType, CFDataRef inAddress, const void* inData, void* inContext );

#pragma mark -

@interface CantripSocket (Private)

- (id)initWithNativeSocket:(CFSocketNativeHandle)inNativeSocket remoteAddress:(struct sockaddr_in*)inRemoteAddress;

- (BOOL)socketCreate;
- (void)socketConnected;
- (void)socketConnectTimeout:(NSTimer*)inTimer;
- (void)socketDisconnected;
- (void)socketNewConnection;
- (void)socketReadable;
- (void)socketWritable;
- (void)socketSend;
- (void)socketRead;
- (CantripSocket*)socketAccept;
- (int)socketAvailableByteCount;
- (BOOL)nonblockingError;

@end

#pragma mark -

@implementation CantripSocket

- (id)init
{
	if( ![super init] )
		return nil;
	
	// Setup socket
	mSocketRef = NULL;
	
	// Setup read and write buffers
	mWritableData = nil;
	mReadableData = nil;
	
	// Setup connect timer
	mConnectTimer = nil;
	
	// Setup host properties
	mRemotePort = NetSocketInvalidPort;
	mRemoteHost = nil;
	
	// Setup state flags
	mConnected = NO;
	mListening = NO;
	mWritable = NO;
	
	return self;
}

- (id)initWithNativeSocket:(CFSocketNativeHandle)inNativeSocket remoteAddress:(struct sockaddr_in*)inRemoteAddress
{
	CFSocketContext socketContext;
	
	if( ![super init] )
		return nil;
	
	// Setup socket context
	bzero( &socketContext, sizeof( CFSocketContext ) );
	socketContext.info = self;
	
	// Create socket based on native socket
	mSocketRef = CFSocketCreateWithNative( kCFAllocatorDefault, inNativeSocket, kCFSocketConnectCallBack + kCFSocketReadCallBack + kCFSocketWriteCallBack, &callbackSocket, &socketContext );
	
	// Setup read and write buffers
	mWritableData = nil;
	mReadableData = nil;
	
	// Setup connect timer
	mConnectTimer = nil;
	
	// Setup host properties
	mRemotePort = inRemoteAddress->sin_port;
	mRemoteHost = [[CantripSocket stringWithAddress:&inRemoteAddress->sin_addr] retain];
	
	// Setup state flags
	mConnected = YES;
	mListening = NO;
	mWritable = YES;
	
	// If creation of the socket failed, let us release ourselves and return nil
	if( mSocketRef == NULL )
	{
		NSLog( @"NetSocket::Failed to create socket" );
		[self release];
		return nil;
	}
	
	return self;
}

- (void)dealloc
{
	// Close socket
	[self close];
	
	// Release host properties
	[mRemoteHost release];
	mRemoteHost = nil;
	mRemotePort = NetSocketInvalidPort;
	
	// Release write buffer
	[mWritableData release];
	mWritableData = nil;
	
	// Release read buffer
	[mReadableData release];
	mReadableData = nil;
	
	// Forget about our delegate
	mDelegate = nil;
	
	[super dealloc];
}

#pragma mark -

+ (CantripSocket*)socket
{
	CantripSocket* socket = [[[CantripSocket alloc] init] autorelease];
	
	if( ![socket open] )
		return nil;
	
	return socket;
}

+ (CantripSocket*)socketListeningOnPort:(UInt16)inPort
{
	CantripSocket* socket = [CantripSocket socket];
	
	if( ![socket listenOnPort:inPort] )
		return nil;
	
	return socket;
}

+ (CantripSocket*)socketConnectedToHost:(NSString*)inHost port:(UInt16)inPort
{
	CantripSocket* socket = [CantripSocket socket];
	
	if( ![socket connectToHost:inHost port:inPort] )
		return nil;
	
	return socket;
}

+ (NSString*)stringWithAddress:(struct in_addr*)inAddress
{
	return [NSString stringWithCString:inet_ntoa( *inAddress )];
}

#pragma mark -

- (id)delegate
{
	return mDelegate;
}

- (void)setDelegate:(id)inDelegate
{
	mDelegate = inDelegate;
}

#pragma mark -

- (BOOL)open
{
	BOOL						result;
	CFRunLoopSourceRef	runloopSource;
	
	// Create socket if necessary
	if( !mSocketRef )
	{
		result = [self socketCreate];
		if( !result )
			return NO;
	}
	
	// Create socket runloop source
	runloopSource = CFSocketCreateRunLoopSource( kCFAllocatorDefault, mSocketRef, 0 );
	if( runloopSource == NULL )
	{
		NSLog( @"NetSocket::Failed to create runloop source" );
		return NO;
	}
	
	// Schedule socket runloop source with current runloop and release
	CFRunLoopAddSource( [[NSRunLoop currentRunLoop] getCFRunLoop], runloopSource, kCFRunLoopDefaultMode );
	CFRelease( runloopSource );
	
	return YES;
}

- (void)close
{
	// Remove socket from runloop and release if necessary
	if( mSocketRef )
	{
		CFSocketInvalidate( mSocketRef );
		CFRelease( mSocketRef );
		mSocketRef = NULL;
	}
	
	// Invalidate connection timer and release
	[mConnectTimer invalidate];
	[mConnectTimer release];
	mConnectTimer = nil;
	
	// Reset state flags
	mConnected = NO;
	mListening = NO;
	mWritable = NO;
	
	if( ![self opened] )
		return;
}

#pragma mark -

- (BOOL)connectToHost:(NSString*)inHost port:(UInt16)inPort
{
	return [self connectToHost:inHost port:inPort timeout:-1.0];
}

- (BOOL)connectToHost:(NSString*)inHost port:(UInt16)inPort timeout:(NSTimeInterval)inTimeout
{
	CFDataRef				socketAddressData;
	struct sockaddr_in	socketAddress;
	struct hostent*		socketHost;
	CFSocketError			socketError;
	
	// Get remote host information (Anyone know if gethostbyname is threadsafe under OSX?)
	socketHost = gethostbyname( [inHost cString] );
	if( socketHost == NULL )
		return NO;
	
	// Setup socket address
	bzero( &socketAddress, sizeof( struct sockaddr_in ) );
	bcopy( (char*)socketHost->h_addr, (char*)&socketAddress.sin_addr, socketHost->h_length );
	socketAddress.sin_family = PF_INET;
	socketAddress.sin_port = htons( inPort );
	
	// Enclose socket address in a data object
	socketAddressData = (CFDataRef)[NSData dataWithBytes:(UInt8*)&socketAddress length:sizeof( struct sockaddr_in )];
		
	// Connect socket
	socketError = CFSocketConnectToAddress( mSocketRef, socketAddressData, -1.0 );
	if( socketError != kCFSocketSuccess )
		return NO;
	
	// Schedule timeout timer if requested
	if( inTimeout >= 0.0 )
		mConnectTimer = [[NSTimer scheduledTimerWithTimeInterval:inTimeout target:self selector:@selector( socketConnectTimeout: ) userInfo:nil repeats:NO] retain];
	
	// Release host properties
	[mRemoteHost release];
	mRemoteHost = nil;
	mRemotePort = NetSocketInvalidPort;
	
	// Setup host properties
	mRemoteHost = [inHost copy];
	mRemotePort = inPort;
	
	return YES;
}

- (BOOL)listenOnPort:(UInt16)inPort
{
	return [self listenOnPort:inPort maxPendingConnections:NetSocketMaxPendingConnections];
}

- (BOOL)listenOnPort:(UInt16)inPort maxPendingConnections:(int)inMaxPendingConnections
{
	CFSocketNativeHandle	nativeSocket;
	struct sockaddr_in	socketAddress;
	int						flag = 1;
	int						result;
	
	// Get native socket
	nativeSocket = CFSocketGetNative( mSocketRef );
	
	// Setup socket options
	result = setsockopt( nativeSocket, SOL_SOCKET, SO_REUSEADDR, &flag, sizeof( flag ) );
	if( result < 0 )
		return NO;
	
	// Setup socket address
	bzero( &socketAddress, sizeof( struct sockaddr_in ) );
	socketAddress.sin_family = PF_INET;
	socketAddress.sin_addr.s_addr = htonl( INADDR_ANY );
	socketAddress.sin_port = htons( inPort );
	
	// Bind socket to port
	result = bind( nativeSocket, (struct sockaddr*)&socketAddress, sizeof( struct sockaddr_in ) );
	if( result < 0 )
		return NO;
	
	// Give socket some ears -- big motherfuckers, Dumbo style
	result = listen( nativeSocket, inMaxPendingConnections );
	if( result < 0 )
		return NO;
	
	mListening = YES;

	return YES;
}

#pragma mark -

- (unsigned)read:(void*)inBuffer amount:(unsigned)inAmount
{
	unsigned amountRead;
	
	// Return if there is no data to read
	if( !mReadableData || [mReadableData length] == 0 )
		return 0;
	
	// Determine how much to actually read
	amountRead = MIN( inAmount, [mReadableData length] );
	
	// Rip out bytes
	[mReadableData getBytes:inBuffer length:amountRead];
	[mReadableData replaceBytesInRange:NSMakeRange( 0, amountRead ) withBytes:NULL length:0];
	
	return amountRead;
}

- (unsigned)readOntoData:(NSMutableData*)inData
{
	unsigned amountRead;
	
	// Return if there is no data to read
	if( !mReadableData || [mReadableData length] == 0 )
		return 0;
	
	// Remember the size of our read buffer
	amountRead = [mReadableData length];
	
	// Rip out bytes
	[inData appendData:mReadableData];
	[mReadableData setLength:0];
	
	return amountRead;
}

- (unsigned)readOntoData:(NSMutableData*)inData amount:(unsigned)inAmount
{
	unsigned amountRead;
	
	// Return if there is no data to read
	if( !mReadableData || [mReadableData length] == 0 )
		return 0;
	
	// Determine how much to actually read
	amountRead = MIN( inAmount, [mReadableData length] );
	
	// Rip out bytes
	[inData appendBytes:[mReadableData bytes] length:amountRead];
	[mReadableData replaceBytesInRange:NSMakeRange( 0, amountRead ) withBytes:NULL length:0];
	
	return amountRead;
}

- (unsigned)readOntoString:(NSMutableString*)inString encoding:(NSStringEncoding)inEncoding amount:(unsigned)inAmount
{
	NSData*		readData;
	NSString*	readString;
	unsigned 	amountRead;
	
	// Return if there is no data to read
	if( !mReadableData || [mReadableData length] == 0 )
		return 0;
	
	// Determine how much to actually read
	amountRead = MIN( inAmount, [mReadableData length] );
	
	// Reference our read buffer, we cut down a little overhead when the the requested amount is the entire buffer.
	if( amountRead == [mReadableData length] )
		readData = mReadableData;
	else
		readData = [[NSData alloc] initWithBytesNoCopy:(void*)[mReadableData bytes] length:amountRead freeWhenDone:NO];
	
	// Create string from bytes with specified encoding
	readString = [[NSString alloc] initWithData:readData encoding:inEncoding];
	
	// Rip out bytes
	[mReadableData replaceBytesInRange:NSMakeRange( 0, amountRead ) withBytes:NULL length:0];
	
	// Append string
	[inString appendString:readString];
	
	// Release temporary string
	[readString release];
	
	// Release temporary read buffer if necessary
	if( readData != mReadableData )
		[readData release];
	
	return amountRead;
}

- (NSData*)readData
{
	NSData* readData;
	
	// Return an empty data object if there are no bytes to read
	if( !mReadableData || [mReadableData length] == 0 )
		return [NSData data];
	
	// Create new data object with contents of our read buffer
	readData = [NSData dataWithData:mReadableData];
	
	// Size our read buffer
	[mReadableData setLength:0];
	
	return readData;
}

- (NSData*)readData:(unsigned)inAmount
{
	NSData*	readData;
	unsigned	amountRead;
	
	// Return an empty data object if there are no bytes to read
	if( !mReadableData || [mReadableData length] == 0 )
		return [NSData data];
	
	// Determine how much to actually read
	amountRead = MIN( inAmount, [mReadableData length] );
	
	// Rip out bytes
	readData = [NSData dataWithBytes:[mReadableData bytes] length:amountRead];
	[mReadableData replaceBytesInRange:NSMakeRange( 0, amountRead ) withBytes:NULL length:0];
	
	return readData;
}

- (NSString*)readString:(NSStringEncoding)inEncoding
{
	NSString* readString;
	
	// Return an empty string object if there are no bytes to read
	if( !mReadableData || [mReadableData length] == 0 )
		return [NSString string];
	
	// Rip out bytes
	readString = [[[NSString alloc] initWithData:mReadableData encoding:inEncoding] autorelease];
	[mReadableData setLength:0];
	
	return readString;
}

- (NSString*)readString:(NSStringEncoding)inEncoding amount:(unsigned)inAmount
{
	NSString*	readString;
	NSData*		readData;
	unsigned		amountRead;
	
	// Return an empty string object if there are no bytes to read
	if( !mReadableData || [mReadableData length] == 0 )
		return [NSString string];
	
	// Determine how much to actually read
	amountRead = MIN( inAmount, [mReadableData length] );
	
	// Reference our read buffer, we cut down a little overhead when the the requested amount is the entire buffer.
	if( amountRead == [mReadableData length] )
		readData = mReadableData;
	else
		readData = [[NSData alloc] initWithBytesNoCopy:(void*)[mReadableData bytes] length:amountRead freeWhenDone:NO];
	
	// Create string from bytes with specified encoding
	readString = [[[NSString alloc] initWithData:readData encoding:inEncoding] autorelease];
	
	// Rip out bytes
	[mReadableData replaceBytesInRange:NSMakeRange( 0, amountRead ) withBytes:NULL length:0];
	
	// Release temporary read buffer if necessary
	if( readData != mReadableData )
		[readData release];
	
	return readString;
}

#pragma mark -

- (void)write:(const void*)inBytes amount:(unsigned)inAmount
{
	// Create write buffer if necessary
	if( !mWritableData )
		mWritableData = [[NSMutableData alloc] init];
	
	// Append data to our write buffer
	[mWritableData appendBytes:inBytes length:inAmount];
	
	// Send as much data as possible
	[self socketSend];
}

- (void)writeData:(NSData*)inData
{
	if( [inData length] == 0 )
		return;
	
	[self write:[inData bytes] amount:[inData length]];
}

- (void)writeString:(NSString*)inString encoding:(NSStringEncoding)inEncoding
{
	if( [inString length] == 0 )
		return;
	
	[self writeData:[inString dataUsingEncoding:inEncoding]];
}

#pragma mark -

- (const NSData*)peekData
{
	NSData* peekedData;
	
	// Return an empty data object if there are no bytes to read
	if( !mReadableData || [mReadableData length] == 0 )
		return [NSData data];
	
	// Set the peeked data to our readable buffer
	peekedData = mReadableData;
	
	return peekedData;
}

#pragma mark -

- (const CFSocketRef)cfsocket
{
	return mSocketRef;
}

- (NSString*)remoteHost
{
	return mRemoteHost;
}

- (UInt16)remotePort
{
	return mRemotePort;
}

- (BOOL)opened
{
	BOOL result;
	
	result = ( mSocketRef != NULL );
	if( !result )
		return NO;
	
	return CFSocketIsValid( mSocketRef );
}

- (BOOL)connected
{
	return mConnected;
}

- (BOOL)listening
{
	return mListening;
}

- (BOOL)readable
{
	return ( [mReadableData length] > 0 );
}

- (unsigned)readableByteCount
{
	return [mReadableData length];
}

- (BOOL)nonblockingError
{
	BOOL result;
	
	switch( errno )
	{
		case EWOULDBLOCK:
		case EALREADY:
		case EINPROGRESS:
			result = YES;
			break;
		
		default:
			result = NO;
			break;
	}
	
	return result;
}

#pragma mark -

- (BOOL)socketCreate
{
	CFSocketContext socketContext;
	
	// Return if the socket has already been created
	if( mSocketRef )
		return NO;
	
	// Setup socket context
	bzero( &socketContext, sizeof( CFSocketContext ) );
	socketContext.info = self;
	
	// Create socket
	mSocketRef = CFSocketCreate( kCFAllocatorDefault, PF_INET, SOCK_STREAM, IPPROTO_TCP, kCFSocketConnectCallBack + kCFSocketReadCallBack + kCFSocketWriteCallBack, &callbackSocket, &socketContext );
	if( mSocketRef == NULL )
	{
		NSLog( @"NetSocket::Failed to create socket" );
		return NO;
	}
	
	return YES;
}

- (void)socketConnected
{
	// Note that we are connected
	mConnected = YES;
	
	// Invalidate connection timer and release
	[mConnectTimer invalidate];
	[mConnectTimer release];
	mConnectTimer = nil;
	
	// Reset our read buffer if necessary.
	// We do it here because the socket may disconnect before the user has read the contents.
	// If the user has connected to a new endpoint, we assume the contents are now no longer relavent.
	if( mReadableData )
		[mReadableData setLength:0];
	
	// Notify delegate
	if( [mDelegate respondsToSelector:@selector( socketConnected: )] )
		[mDelegate socketConnected:self];
	
	// Try sending any data that was added to our write buffer while the socket was connecting
	[self socketSend];
}

- (void)socketConnectTimeout:(NSTimer*)inTimer
{
	NSTimeInterval timeInterval;
	
	// Remember timer interval
	timeInterval = [mConnectTimer timeInterval];

	// Close socket
	[self close];
		
	// Notify delegate
	if( [mDelegate respondsToSelector:@selector( socket:connectionTimedOut: )] )
		[mDelegate socket:self connectionTimedOut:timeInterval];
}

- (void)socketDisconnected
{
	// Note that we are disconnected
	mConnected = NO;
	
	// Reset our write buffer if necessary
	if( mWritableData )
		[mWritableData setLength:0];
		
	// Close socket
	[self close];
	
	// Notify delegate
	if( [mDelegate respondsToSelector:@selector( socketDisconnected: )] )
		[mDelegate socketDisconnected:self];
}

- (void)socketNewConnection
{
	CantripSocket* socket;
	
	// Accept connection
	socket = [self socketAccept];
	if( !socket )
		return;
	
	// Notify delegate
	if( [mDelegate respondsToSelector:@selector( socket:newConnection: )] )
		[mDelegate socket:self newConnection:socket];
}

- (void)socketReadable
{
	// Read in available data
	[self socketRead];
	
	// This read actually disconnected us, so just return
	if( !mConnected )
		return;
	
	// Notify delegate
	if( [mDelegate respondsToSelector:@selector( socket:dataAvailable: )] )
		[mDelegate socket:self dataAvailable:[mReadableData length]];
}

- (void)socketWritable
{
	// Note that we are writable
	mWritable = YES;
	
	// Try sending any more data from our buffer
	[self socketSend];
}

- (void)socketSend
{
	CFSocketNativeHandle	nativeSocket;
	int						amountSent;
	
	// Return if the socket is not writable or hasn't been created
	if( !mWritable )
		return;
	
	// Return if there is no data to be written
	if( [mWritableData length] == 0 )
		return;
	
	// Get the native socket
	nativeSocket = CFSocketGetNative( mSocketRef );
	
	// Queue data for sending
	amountSent = send( nativeSocket, [mWritableData bytes], [mWritableData length], 0 );
	
	// Check our result
	if( amountSent == [mWritableData length] )
		mWritable = YES;
	else
	if( amountSent >= 0 || [self nonblockingError] )
		mWritable = NO;
	else
	{
		[self close];
		return;
	}
	
	// Remove sent data from our buffer
	[mWritableData replaceBytesInRange:NSMakeRange( 0, amountSent ) withBytes:NULL length:0];
	
	// Notify delegate when our write buffer has been exhausted
	if( [mWritableData length] == 0 )
		if( [mDelegate respondsToSelector:@selector( socketSentData: )] )
			[mDelegate socketSentData:self];
}

- (void)socketRead
{
	CFSocketNativeHandle	nativeSocket;
	void*						readBuffer;
	int						amountAvailable;
	int						amountRead;
	
	// Determine how much data is actually available on the socket
	amountAvailable = [self socketAvailableByteCount];
	if( amountAvailable == 0 )
	{
		[self socketDisconnected];
		return;
	}
	else
	if( amountAvailable < 0 )
		return;
	
	// Create temporary read buffer
	readBuffer = malloc( amountAvailable );
	if( !readBuffer )
		return;
	
	// Get the native socket
	nativeSocket = CFSocketGetNative( mSocketRef );
	
	// Read available bytes off the socket
	amountRead = recv( nativeSocket, readBuffer, amountAvailable, 0 );
	
	// Create read buffer if necessary
	if( !mReadableData )
		mReadableData = [[NSMutableData alloc] init];
	
	// Append data to our write buffer
	[mReadableData appendBytes:readBuffer length:amountRead];
	
	// Free our temporary read buffer
	free( readBuffer );
}

- (CantripSocket*)socketAccept
{
	struct sockaddr_in	socketAddress;
	int						socketAddressSize;
	int						socket;
	int						nativeSocket;
	
	// Get the native socket
	nativeSocket = CFSocketGetNative( mSocketRef );
	
	// Accept incoming connection
	socket = accept( nativeSocket, (struct sockaddr*)&socketAddress, &socketAddressSize );
	if( socket < 0 )
		return nil;
	
	return [[[CantripSocket alloc] initWithNativeSocket:socket remoteAddress:&socketAddress] autorelease];
}

- (int)socketAvailableByteCount
{
	CFSocketNativeHandle	nativeSocket;
	int						bytesAvailable;
	
	nativeSocket = CFSocketGetNative( mSocketRef );
	if( ioctl( nativeSocket, FIONREAD, &bytesAvailable ) != -1 )
		return bytesAvailable;
	else
	{
		if( errno == EINVAL )
			return -1;
		else
			return 0;
	}
}

#pragma mark -

void 
callbackSocket( CFSocketRef inSocketRef, CFSocketCallBackType inType, CFDataRef inAddress, const void* inData, void* inContext )
{
	CantripSocket* socket = (CantripSocket*)inContext;
	if( !socket )
		return;
	
	switch( inType )
	{
		case kCFSocketConnectCallBack:
			[socket socketConnected];
			break;
		
		case kCFSocketReadCallBack:
			{
				if( [socket listening] )
					[socket socketNewConnection];
				else
					[socket socketReadable];
			}
			break;
		
		case kCFSocketWriteCallBack:
			[socket socketWritable];
			break;
		
		default:
			// Alien callback type. I'm talkin' like freaky shit from 'Signs' yo!
			break;
	}
}

@end
